import os
import asyncio
from pathlib import Path
from rich.console import Console
from typing_extensions import Annotated
from dotenv import load_dotenv

import typer
from load.checkpoint import Checkpoint
from load.custom import DocImport

console = Console()

app = typer.Typer(no_args_is_help=True)

@app.command()
def clean(
    env_file: Path = typer.Option(".env", help="Path to the .env file containing environment variables.")
):
    if env_file:
        load_dotenv(env_file, override=True)

    cp = Checkpoint()
    cp.reset()

@app.command()
def process(
    override: Annotated[bool, typer.Option(prompt=True, help="覆盖已经生成的文件么?")] = False,
    source: Annotated[str, typer.Option(prompt=True, help="原始文件(.txt) 路径")] = "./tests/data/input",
    working_dir: Annotated[str, typer.Option(prompt=True, help="工作文件夹")] = "./tests/data/output",
    env_file: Path = typer.Option(".env", help="Path to the .env file containing environment variables.")
):
    """
    Make the indexing.
    """
    if env_file:
        load_dotenv(env_file, override=True)

    if not os.path.exists(source):
        typer.echo(f"Path: {source} Not found")
        return 

    if (override):
        import shutil
        if os.path.exists(working_dir):
            typer.echo(f"Will remove the {working_dir}")
            shutil.rmtree(working_dir)
            # 扫描指定文件下的所有txt文件
    
    os.makedirs(working_dir, mode=511, exist_ok=True)

    d = DocImport(source, working_dir)
    correct, incorrect = asyncio.run(d.process())
    typer.echo(f"Importing {correct} correct News and {incorrect} news")  

    cp = Checkpoint(working_dir)
    asyncio.run(cp.process())

    typer.echo("Done")    

@app.command()
def load(
    working_dir: Annotated[str, typer.Option(prompt=True, help="工作文件夹")] = "./tests/data/output",
    env_file: Path = typer.Option(".env", help="Path to the .env file containing environment variables.")
):
    """
    Make the indexing.
    """
    if env_file:
        load_dotenv(env_file, override=True)

    cp = Checkpoint(working_dir)
    count_nodes, count_descriptions = asyncio.run(cp.batch_import())

    typer.echo(f"Loaded {count_nodes} nodes, and {count_descriptions} descriptions")  

@app.command()
def optimize(
    degree: Annotated[int, typer.Option(prompt=True, help="仅处理超过这个描述数量的实体对象，并对这些对象进行总结")] = 10,
    env_file: Path = typer.Option(".env", help="Path to the .env file containing environment variables.")
):
    """
    Make the indexing.
    """
    if env_file:
        load_dotenv(env_file, override=True)

    cp = Checkpoint()

    count = cp.optimize(degree)

    typer.echo(f"Optimize {count} nodes")  

def main():
    app()

if __name__ == "__main__":
    main()
